# -*- coding: UTF-8 -*-

import time
from selenium import webdriver
import pytesseract
from PIL import Image, ImageEnhance


# 配置环境变量
pytesseract.pytesseract.tesseract_cmd = 'C:/Program Files (x86)/Tesseract-OCR/tesseract.exe'
tessdata_dir_config = '--tessdata-dir "C:/Program Files (x86)/Tesseract-OCR/tessdata"'


browser = webdriver.Chrome()
browser.get(
    "file:///C:/sethnie/Project/WebCrawler/GetResultsOfPostgraduateEntranceExamination/%E5%85%A8%E5%9B%BD%E7%A1%95%E5%A3%AB%E7%A0%94%E7%A9%B6%E7%94%9F%E7%BB%9F%E4%B8%80%E6%8B%9B%E7%94%9F%E8%80%83%E8%AF%95%E9%83%A8%E5%88%86%E8%80%83%E7%94%9F%E5%88%9D%E8%AF%95%E6%88%90%E7%BB%A9%E6%9F%A5%E8%AF%A2%E7%B3%BB%E7%BB%9F.html")

print("[查询时间：" + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + "][查询成功]")


# 窗口最大化
browser.maximize_window()
# 获取图像验证码位置、大小
imgelement = browser.find_element_by_xpath('//*[@id="app"]/div/form[2]/div[5]/div/div[2]/img')
location = imgelement.location
size = imgelement.size
rangle = (
    int(location['x']), int(location['y']), int(location['x'] + size['width']),
    int(location['y'] + size['height'])
)

# 截取网页
browser.get_screenshot_as_file('./img/image1.png')  # 比较好理解
# 等待保存
time.sleep(1)
# 重新裁剪，获取验证码
im = Image.open('./img/image1.png')
region = im.crop(rangle)  # 此时，region是一个新的图像对象
region.save("./img/code1.png")

time.sleep(1)

# 验证码二值化、对比度加强
im = Image.open("./img/code1.png")
imgry = im.convert('L')  # 图像加强，二值化
sharpness = ImageEnhance.Contrast(imgry)  # 对比度增强
sharp_img = sharpness.enhance(2.0)
sharp_img.save("./img/code2.png")

# 去除验证码多余空格
code = pytesseract.image_to_string(sharp_img).replace( ' ' , '' )
print(code)

browser.quit()

